import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from tqdm import tqdm
import torch.nn.functional as F
import matplotlib.pyplot as plt

# 数据预处理
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])

# 加载 CIFAR-10 数据集
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=64, shuffle=True, num_workers=2)

testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=64, shuffle=False, num_workers=2)

# 定义模型
class FuzzyInferenceBlock(nn.Module):
    def __init__(self, S, T, n_feature):
        super(FuzzyInferenceBlock, self).__init__()
        self.S = S  # 模糊分级个数
        self.T = T  # 规则生成层、归一化层节点数
        self.c1 = nn.Parameter(torch.rand(S, )  # 初始化 c1 为可学习参数
        self.b1 = nn.Parameter(torch.rand(S, )  # 初始化 b1 为可学习参数
        self.w1 = nn.Parameter(torch.rand(1, T))  # 初始化 w1 为可学习参数

    def forward(self, inputs):
        Y = inputs.clone()
        u1 = torch.zeros((self.S, inputs.shape[1]), device=inputs.device)
        alpha = []

        # 计算每个输入样本的隶属度
        for m in range(inputs.shape[0]):
            for i in range(self.S):
                for j in range(inputs.shape[1]):
                    u1[i][j] = torch.exp((-1) * ((inputs[m, j] - self.c1[i][j]) ** 2) / (self.b1[i][j] ** 2))

            # 规则生成层计算输出
            alpha1 = torch.zeros((self.T, 1), device=inputs.device)
            for i in range(self.S):
                a = 1
                if i == 1:
                    a = 8
                for p in range(self.S):
                    b = 1
                    if p == 1:
                        b = 4
                    for q in range(self.S):
                        c = 1
                        if q == 1:
                            c = 2
                        for k in range(self.S):
                            alpha1[i * a + p * b + q * c + k] = u1[i][0] * u1[p][1] * u1[q][2] * u1[k][3]

            alpha.append(alpha1)

            # 归一化层计算输出
            alphasum = torch.sum(alpha1)
            # 输出层计算输出
            a2 = torch.dot(self.w1, alpha1)
            Y[m] = a2 / alphasum

        return Y


class FCNN(nn.Module):
    def __init__(self, n_femap=100, stride=2, mu=3.0, sigma=1.2, dropout=True, num_classes=10):
        super(FCNN, self).__init__()
        self.dropout = dropout
        self.conv1 = nn.Conv2d(3, 20, kernel_size=3, stride=2, padding=1)
        self.conv2 = nn.Conv2d(20, 40, kernel_size=3, stride=2, padding=1)
        self.conv3 = nn.Conv2d(40, 40, kernel_size=3, stride=2, padding=1)
        self.conv4 = nn.Conv2d(40, 40, kernel_size=3, stride=1, padding=1)
        self.conv5 = nn.Conv2d(40, n_femap, kernel_size=3, stride=stride, padding=0)

        if self.dropout:
            self.dropout_layer = nn.Dropout(0.2)

        self.n_femap = n_femap
        self.fuzzy_blocks = nn.ModuleList([
            FuzzyInferenceBlock(S=2, T=16, n_feature=n_femap)  # 传入 S 和 T
            for i in range(n_femap)
        ])
        self.fc = nn.Linear(n_femap * 10, num_classes)

        self.fc1 = None

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.max_pool2d(x, kernel_size=3, stride=1, padding=1)

        x = F.relu(self.conv2(x))
        x = F.max_pool2d(x, kernel_size=2, stride=1, padding=1)

        x = F.relu(self.conv4(x))
        x = F.max_pool2d(x, kernel_size=2, stride=1, padding=1)

        x = F.relu(self.conv5(x))
        if self.dropout:
            x = self.dropout_layer(x)

        x = x.view(x.size(0), -1)
        self.fc1 = nn.Linear(x.size(1), 128)  # 根据特征图大小动态创建
        x = F.relu(self.fc1(x))
        fuzzy_outputs = [self.fuzzy_blocks[i](x) for i in range(self.n_femap)]

        merged = torch.cat(fuzzy_outputs, dim=1)

        out = self.fc(merged)
        return out

# 训练设置
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # 检查GPU是否可用
device = torch.device("cpu")
model = FCNN().to(device)  # 移动模型到GPU
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
losses = []

# 训练过程
num_epochs = 100  # 可以根据需要调整
for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for inputs, labels in tqdm(trainloader):
        inputs, labels = inputs.to(device), labels.to(device)  # 移动输入和标签到GPU

        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward(retain_graph=True)
        optimizer.step()

        running_loss += loss.item()

    print(f"Epoch [{epoch + 1}/{num_epochs}], Loss: {running_loss / len(trainloader):.4f}")
    losses.append(running_loss / len(trainloader))

# 绘制损失图
plt.plot(losses)
plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.title('Training Loss Over Epochs')
plt.show()

# 测试过程
model.eval()
correct = 0
total = 0
with torch.no_grad():
    for inputs, labels in testloader:
        inputs, labels = inputs.to(device), labels.to(device)  # 移动输入和标签到GPU
        outputs = model(inputs)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'Accuracy of the model on the 10000 test images: {100 * correct / total:.2f}%')
